# plot_3d_path.py

import json
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import os

def load_paths(path_file='agent_paths.json'):
    """
    加载代理路径数据。

    参数：
    - path_file (str): 存储代理路径的JSON文件路径。

    返回：
    - paths (list): 包含每个Episode路径的列表。
    """
    if os.path.exists(path_file):
        with open(path_file, 'r') as f:
            paths = json.load(f)
        print(f"Loaded {len(paths)} episodes from {path_file}.")
        return paths
    else:
        print(f"No agent paths found at {path_file}.")
        return []

def plot_single_episode_path(path, episode_num, save_dir='episode_plots'):
    """
    绘制单个Episode的3D路径图，并保存为图片文件。

    参数：
    - path (list): 当前Episode的路径列表，包含多个位置点，每个位置点是一个包含三个索引的列表。
    - episode_num (int): 当前Episode的编号。
    - save_dir (str): 保存图片文件的目录。

    返回：
    - None
    """
    # 创建保存目录（如果不存在）
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    # 创建3D图形
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111, projection='3d')

    # 确保路径中有多个点
    if len(path) < 2:
        print(f"Episode {episode_num} 的路径点数不足，跳过绘制。")
        return

    # 解包路径点
    x, y, z = zip(*path)

    # 绘制路径
    ax.plot(x, y, z, marker='o', linestyle='-', color='b', alpha=0.7)

    # 设置轴标签和标题
    ax.set_xlabel('Learning Rate Index (X)')
    ax.set_ylabel('Max Depth Index (Y)')
    ax.set_zlabel('N Estimators Index (Z)')
    ax.set_title(f'DQN Agent Path - Episode {episode_num}')

    # 保存图形到文件
    save_path = os.path.join(save_dir, f'episode_{episode_num}_path.png')
    plt.savefig(save_path)
    print(f"Saved Episode {episode_num} path to {save_path}.")

    # 关闭图形以释放内存
    plt.close(fig)

def visualize_paths_separately(path_file='agent_paths.json', num_episodes=10, save_dir='episode_plots'):
    """
    分别绘制多个Episode的3D路径图，每个Episode一张图。

    参数：
    - path_file (str): 存储代理路径的JSON文件路径。
    - num_episodes (int): 要可视化的Episode数量。
    - save_dir (str): 保存图片文件的目录。

    返回：
    - None
    """
    # 加载路径数据
    paths = load_paths(path_file)

    # 确保有路径数据
    if not paths:
        print(f"No paths to plot from {path_file}.")
        return

    # 确定实际要绘制的Episode数量
    episodes_to_plot = min(num_episodes, len(paths))
    print(f"绘制前 {episodes_to_plot} 个Episode的路径。")

    # 逐个绘制并保存每个Episode的路径
    for i in range(episodes_to_plot):
        path = paths[i]
        episode_num = i + 1  # Episode编号从1开始
        plot_single_episode_path(path, episode_num, save_dir)

    print(f"所有Episode的路径图已保存到目录：{save_dir}")

if __name__ == "__main__":
    # 调用函数，绘制前10个Episode的路径，并保存到'episode_plots'目录
    visualize_paths_separately('agent_paths.json', num_episodes=500, save_dir='episode_plots')
